#include "ICP.h"

#include <QElapsedTimer>
#include <fstream>
#include <cfloat>
#include <limits>
#include <iomanip>


#include "RigidMotionEstimator.h"


ICP::ICP( SurfaceMeshModel* m1, SurfaceMeshModel* m2 )
{
	this->m1 = m1;
	this->m2 = m2;

	// set parameters
	int it = 1;
	double maxdist = 0.1;
	double stop = 1e-5;
	bool is_point_to_plane = true;

    // kd-tree of target data
	target_kdtree = new NanoflannKdTree3d();
	Vector3VertexProperty m2_points = m2->vertex_coordinates();
	foreach (Vertex v, m2->vertices()) target_kdtree->addPoint(m2_points[v]);
	target_kdtree->build();
}

ICP::~ICP()
{
	delete target_kdtree;
}

void ICP::execute()
{
	// convert data to Eigen
	Eigen::Map<Eigen::Matrix3Xd> X((double *)(m1->vertex_coordinates().data()), 3, m1->n_vertices());
	Eigen::Map<Eigen::Matrix3Xd> dst((double *)(m2->vertex_coordinates().data()), 3, m2->n_vertices());
	Eigen::Map<Eigen::Matrix3Xd> dst_normals((double *)(m2->vertex_normals().data()), 3, m2->n_vertices());

	//////////////////////////////////////////////////
	Eigen::Matrix3Xd Yp = Eigen::Matrix3Xd::Zero(3, m1->n_vertices());
	Eigen::Matrix3Xd Yn = Eigen::Matrix3Xd::Zero(3, m1->n_vertices());

	//////////////////////////////////////////////////
	for(int icp = 0; icp < it; ++icp) {

		Eigen::Matrix3Xd src_old = X;

		/// CLOSEST POINT (UPDATES Yp, Yn)
		for(int i=0; i<X.cols(); ++i) {
			int idx = target_kdtree->closest( Vec3d( X(0,i), X(1,i), X(2,i) ) );
			Eigen::Vector3d closest_pt = dst.col(idx);
			Eigen::Vector3d closest_nr = dst_normals.col(idx);
			Yp.col(i) = closest_pt;
			Yn.col(i) = closest_nr;
		}


		/// ROTATION ESTIMATION (UPDATES R,t)
		/// Covariance entries weights
		Eigen::VectorXd weights = Eigen::VectorXd::Ones(X.cols());

		/// (SET WEIGHTS WHERE DIST>MAXDIST TO ZERO)
		weights = (X-Yp).colwise().norm();
		weights.array() /= std::min(maxdist, weights.maxCoeff()) + 1e-8;
		weights = weights.array().min(1.0);
		weights.array() *= -1.0;
		weights.array() += 1.0;

		/// Which version of ICP should we be using?
		if(is_point_to_plane) 
			RigidMotionEstimator::point_to_plane(X, Yp, Yn, weights);
		else
			RigidMotionEstimator::point_to_point(X, Yp, weights);

		/// stopping criteria
		double final = (X - src_old).colwise().norm().maxCoeff();
		if(final < stop) break;
	}
}
